import mmcv
import os
import argparse

from mmdet.apis import init_detector, inference_detector, show_result_pyplot

parser = argparse.ArgumentParser(description="args to use this scripts")
parser.add_argument("config_file", type=str)
parser.add_argument("checkpoint_file", type=str)
parser.add_argument("imgdir", type=str)
parser.add_argument("result_dir", type=str)
#  config_file = "/workspace/volume/wangxianzhuo-data/iva_experiment/configs/2021half/"
#  checkpoint_file = (
#  "/workspace/volume/wangweizhuo-data/mmfacecar/result/baseline-offset/epoch_24.pth"
#  )
# build the model from a config file and a checkpoint file
# test a single image
#  imgdir = "/workspace/volume/wangweizhuo-data/mmfacecar/test-img/0--Parade/"
#  resultdir = "/workspace/volume/wangweizhuo-data/mmfacecar/result-img/0--Parade/"

args = parser.parse_args()
model = init_detector(args.config_file, args.checkpoint_file, device="cuda:0")
imgList = os.listdir(args.imgdir)
for count in range(len(imgList)):
    print(count)
    im_name = imgList[count]
    im_path = os.path.join(args.imgdir, im_name)
    # img = '/workspace/volume/wangweizhuo-data/mmfacecar/test-img/0_Parade_marchingband_1_95.jpg'
    result = inference_detector(model, im_path)

    if hasattr(model, "module"):
        model = model.module
    img = model.show_result(
        im_path,
        result,
        score_thr=0.3,
        show=False,
        out_file=args.result_dir + "0.5/" + im_name,
    )
